package com.wind.rag;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

/**
 * @author wind
 * @version 1.0
 * @description: 向量检索示例
 * @date 2025/3/18 18:55
 */
public class vector {

    public static void main(String[] args) {
        // 创建向量模型，使用 BAAI/bge-m3 模型
        // 配置 API 密钥和基础 URL
        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .modelName("BAAI/bge-m3")
                .apiKey("sk-afpfqajtcpqlgnnqicmixykrwlbzrqlimclafeqpnjkotfoe")
                .baseUrl("https://api.siliconflow.cn/v1")
                .build();

        // 将文本"水果是西瓜"转换为向量表示
        Response<Embedding> embedded = embeddingModel.embed("水果是西瓜");

        // 创建内存向量存储，用于存储向量和对应的文本
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        
        // 将向量和对应的文本存入向量存储
        embeddingStore.add(embedded.content(), TextSegment.from("水果是西瓜"));

        // 设置查询文本
        String query = "水果";
        
        // 将查询文本转换为向量
        Response<Embedding> embedding = embeddingModel.embed(query);

        // 执行向量相似度搜索
        // minScore 设置为 0.7，表示只返回相似度大于 0.7 的结果
        EmbeddingSearchResult<TextSegment> searched = embeddingStore.search(EmbeddingSearchRequest.builder()
                .queryEmbedding(embedding.content())
                .minScore(0.7)
                .build());

        // 打印搜索结果中最相似的文本

        System.out.println("打印搜索结果中最相似的文本");
        System.out.println(searched.matches().get(0).embedded().text() +" 相似分数：" +
                searched.matches().get(0).score());
    }
}
